package drds.data_migrate.job;


import drds.common.Author;
import drds.data_migrate.Authors;
import drds.data_migrate.vo.JobConfig;
import drds.data_migrate.vo.TaskConfig;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.Validate;

import java.util.*;

@Author(name = Authors.LI_YANG)
public final class TaskAssigns
{
    /**
     * 公平的分配 wrapper 到对应的 taskGroup 中。
     * 公平体现在：会考虑 wrapper 中对资源负载作的 load 标识进行更均衡的作业分配操作。
     * TODO 具体文档举例说明
     */
    public static List<TaskGroup> resourceLoadBalanceAssign(JobConfig config, int readWriteTaskNumber, int readWriteTaskNumberPerTaskGroup)
    {
        Validate.isTrue(config != null, "框架获得的 job 不能为 null.");
//
        List<TaskConfig> taskTaskConfigList = config.getTaskConfigList();
        Validate.isTrue(taskTaskConfigList.size() > 0, "框架获得的切分后的 job 无内容.");
//
        Validate.isTrue(readWriteTaskNumber > 0 && readWriteTaskNumberPerTaskGroup > 0, "每个channel的平均task数[averTaskPerChannel]，channel数目[channelNumber]，每个taskGroup的平均channel数[channelsPerTaskGroup]都应该为正数");
        int taskGroupNumber = (int) Math.ceil(1.0 * readWriteTaskNumber / readWriteTaskNumberPerTaskGroup);
//
        TaskConfig taskConfig0 = taskTaskConfigList.get(0);
        String readerResourceMark = taskConfig0.getJobReaderParameterLoadBalanceResourceMark();
        String writerResourceMark = taskConfig0.getJobWriterParameterLoadBalanceResourceMark();
        //相同的资源均匀打散
        boolean hasLoadBalanceResourceMark = StringUtils.isNotBlank(readerResourceMark) || StringUtils.isNotBlank(writerResourceMark);
        if (!hasLoadBalanceResourceMark)
        {
            // fake 一个固定的 key 作为资源标识（在 reade 或者 write 上均可，此处选择在 reade 上进行 fake）
            for (TaskConfig taskConfig1 : taskTaskConfigList)
            {
                //Constants.jobReaderParameter + "." + Constants.loadBalanceResourceMark,
                taskConfig1.setJobReader("aFakeResourceMarkForLoadBalance");
            }
            // 是为了避免某些插件没有设置 资源标识 而进行了一次随机打乱操作
            Collections.shuffle(taskTaskConfigList, new Random(System.currentTimeMillis()));
        }
        //
        LinkedHashMap<String, List<Integer>> shuffleResourceMarkToTaskIdListMap = getShuffleResourceMarkToTaskIdListMap(taskTaskConfigList);
        List<TaskGroup> taskGroupTaskConfigList = assign(shuffleResourceMarkToTaskIdListMap, config, taskGroupNumber);
        // 调整 每个 taskGroup 对应的 RecordQueue 个数（属于优化范畴）
        adjustReadWriteTaskNumberPerReadWriteTaskGroup(taskGroupTaskConfigList, readWriteTaskNumber);
        return taskGroupTaskConfigList;
    }

    /**
     * 根据task 配置，获取到：
     * 资源名称 --> taskId(List) 的 map 映射关系
     */
    private static LinkedHashMap<String, List<Integer>> getShuffleResourceMarkToTaskIdListMap(List<TaskConfig> taskConfigList)
    {
        // key: resourceMark, value: taskId
        LinkedHashMap<String, List<Integer>> readerResourceMarkToTaskIdListMap = new LinkedHashMap<String, List<Integer>>();
        LinkedHashMap<String, List<Integer>> writerResourceMarkToTaskIdListMap = new LinkedHashMap<String, List<Integer>>();

        for (TaskConfig taskConfig : taskConfigList)
        {
            int taskId = taskConfig.getTaskId();
            // 把 readerResourceMark 加到 readerResourceMarkAndTaskIdMap 中
            String readerResourceMark = taskConfig.getJobReaderParameterLoadBalanceResourceMark();//(Constants.);
            if (readerResourceMarkToTaskIdListMap.get(readerResourceMark) == null)
            {
                readerResourceMarkToTaskIdListMap.put(readerResourceMark, new LinkedList<Integer>());
            }
            readerResourceMarkToTaskIdListMap.get(readerResourceMark).add(taskId);

            // 把 writerResourceMark 加到 writerResourceMarkAndTaskIdMap 中
            String writerResourceMark = taskConfig.getJobWriterParameterLoadBalanceResourceMark();
            if (writerResourceMarkToTaskIdListMap.get(writerResourceMark) == null)
            {
                writerResourceMarkToTaskIdListMap.put(writerResourceMark, new LinkedList<Integer>());
            }
            writerResourceMarkToTaskIdListMap.get(writerResourceMark).add(taskId);
        }

        if (readerResourceMarkToTaskIdListMap.size() >= writerResourceMarkToTaskIdListMap.size())
        {
            // 采用 reade 对资源做的标记进行 shuffle
            return readerResourceMarkToTaskIdListMap;
        } else
        {
            // 采用 write 对资源做的标记进行 shuffle
            return writerResourceMarkToTaskIdListMap;
        }
    }

    /**
     * /**
     * 需要实现的效果通过例子来说是：
     * <pre>
     * a 库上有表：0, 1, 2
     * a 库上有表：3, 4
     * c 库上有表：5, 6, 7
     *
     * 如果有 4个 taskGroup
     * 则 assign 后的结果为：
     * taskGroup-0: 0,  4,
     * taskGroup-1: 3,  6,
     * taskGroup-2: 5,  2,
     * taskGroup-3: 1,  7
     *
     * </pre>
     */
    private static List<TaskGroup> assign(LinkedHashMap<String, List<Integer>> resourceMarkToTaskIdListMap, JobConfig jobConfig, int taskGroupNumber)
    {
        List<TaskConfig> jobTaskConfigList = jobConfig.getTaskConfigList();
        JobConfig taskGroupTemplate = jobConfig;
        //taskGroupTemplate.remove(Constants.taskConfigList);
        //taskGroup
        List<List<TaskConfig>> taskGroupConfigurationListList = new ArrayList<List<TaskConfig>>(taskGroupNumber);
        for (int i = 0; i < taskGroupNumber; i++)
        {
            taskGroupConfigurationListList.add(new LinkedList<TaskConfig>());
        }

        int maxLength = -1;
        List<String> resourceMarkList = new ArrayList<String>();
        for (Map.Entry<String, List<Integer>> entry : resourceMarkToTaskIdListMap.entrySet())
        {
            resourceMarkList.add(entry.getKey());
            if (entry.getValue().size() > maxLength)
            {
                maxLength = entry.getValue().size();
            }
        }
        //
        int taskGroupIndex = 0;
        for (int i = 0; i < maxLength; i++)
        {
            //依次遍历
            for (String resourceMark : resourceMarkList)
            {
                if (resourceMarkToTaskIdListMap.get(resourceMark).size() > 0)
                {
                    int taskId = resourceMarkToTaskIdListMap.get(resourceMark).get(0);
                    taskGroupConfigurationListList.get(taskGroupIndex % taskGroupNumber).add(jobTaskConfigList.get(taskId));
                    resourceMarkToTaskIdListMap.get(resourceMark).remove(0);
                    //
                    taskGroupIndex++;


                }
            }
        }
        List<TaskGroup> taskConfigList = new LinkedList<TaskGroup>();

        for (int i = 0; i < taskGroupNumber; i++)
        {
            TaskGroup taskConfig = null;
            taskConfig.setTaskGroupId(i);
            taskConfig.setTaskConfigList(taskGroupConfigurationListList.get(i));
            //
            taskConfigList.add(taskConfig);
        }

        return taskConfigList;
    }


    private static void adjustReadWriteTaskNumberPerReadWriteTaskGroup(List<TaskGroup> taskGroupTaskConfigList, int readWriteTaskNumber)
    {
        int taskGroupNumber = taskGroupTaskConfigList.size();
        int readWriteTaskNumberPerTaskGroup = readWriteTaskNumber / taskGroupNumber;
        int remainderCount = readWriteTaskNumber % taskGroupNumber;
        // 表示有 remainderChannelCount 个 taskGroup,其对应 RecordQueue 个数应该为：avgChannelsPerTaskGroup + 1；
        // （taskGroupNumber - remainderChannelCount）个 taskGroup,其对应 RecordQueue 个数应该为：avgChannelsPerTaskGroup
        int i = 0;
        //count
        for (; i < remainderCount; i++)
        {
            taskGroupTaskConfigList.get(i).setReadWriteTaskNumberOfThisTaskGroup(readWriteTaskNumberPerTaskGroup + 1);
        }

        for (int j = 0; j < taskGroupNumber - remainderCount; j++)
        {
            taskGroupTaskConfigList.get(i + j).setReadWriteTaskNumberOfThisTaskGroup(readWriteTaskNumberPerTaskGroup);
        }
    }


}
